#include "codegen.hpp"
#include "ast.hpp"
#include "opcode.hpp"
#include "symtable.hpp"
#include "codeobj.hpp"

#include <cassert>
#include <iostream>

using namespace std;

void Label::resolve(vector<unsigned char>& codes) {
    _pos = (unsigned short)codes.size();
    for (auto it = _uses.begin(); it != _uses.end(); it++) {
        int old_value = (((unsigned int)codes[*it-1]) << 8) 
            | ((unsigned int)codes[*it-2]);

        int param = 0;
        if (old_value == 1) {// absolute jump
            param = _pos;
        }
        else {
            param = _pos - *it;
        }
        codes[*it-1] = (unsigned char)((param >> 8) & 0xff);
        codes[*it-2] = (unsigned char)(param & 0xff);
    }

    _resolved = true;
}

CodeGen::CodeGen(SymbolTable* symtable, ostream& out, bool is_func) : _offset(0), _stack_size(0), _nlocals(0), _arg_count(0),
    _out(out), _indent(0), _is_lambda(false),
    _flags(0), _sym_table(symtable), _is_func(is_func) {
    _context = Context::Load;
    string hilang_home = getenv("HILANG_HOME");
    _out << "import sys\n";
    _out << "sys.path.append(\"" << hilang_home << "/pylib\")\n";
    _out << "from std import *\n";
}

CodeGen::~CodeGen() {
}

void CodeGen::indent() {
    for (int i = 0; i < _indent; i++) {
        _out << "    ";
    }
}

void CodeGen::prelogue(Node* root) {
    _type_checker->visit(root);
    if (_type_checker->has_error()) {
        exit(-1);
    }
}

Node* CodeGen::visit(Node* n) {
    if (n->silent()) {
        return n;
    }

    return n->accept(this);
}

void CodeGen::add_op(unsigned char op_code, unsigned short param) {
    _insts.push_back(op_code);
    _offset += 1;
    if (op_code >= HAVE_ARGUMENT) {
        _insts.push_back((unsigned char)(param & 0xff));
        _insts.push_back((unsigned char)((param >> 8) & 0xff));
        _offset += 2;
    }
}

void CodeGen::bind(Label& label) {
    label.resolve(_insts);
}

TypeTable* CodeGen::type_table() {
    return _sym_table->type_table();
}

Node* CodeGen::visit(PrintNode* n) {
    _out << "print(";
    visit(n->value());
    _out << ")";
    return n;
}

Node* CodeGen::visit(ReturnNode* n) {
    if (!_is_lambda) {
        _out << "return ";
    }
    visit(n->value());
    return n;
}

Node* CodeGen::visit(BinaryOp* op) {
    if (op->op_type() == AST_OP_COMPOSE) {
        _out << "compose(";
        visit(op->left());
        _out << ", ";
        visit(op->right());
        _out << ")";
        return op;
    }

    _out << "(";
    visit(op->left());
    Type* left_type = op->left()->get_type(type_table());
    _out << ")";

    switch(op->op_type()) {
        case AST_OP_ADD:
            _out << " + ";
            break;
        case AST_OP_SUB:
            _out << " - ";
            break;
        case AST_OP_MUL:
            _out << " * ";
            break;
        case AST_OP_DIV:
            if (left_type == IntType::get_instance()) {
                _out << " // ";
            }
            else {
                _out << " / ";
            }
            break;
        case AST_OP_MOD:
            _out << " % ";
            break;
        case AST_OP_BIT_AND:
            _out << " & ";
            break;
        case AST_OP_BIT_OR:
            _out << " | ";
            break;
        case AST_OP_BIT_XOR:
            _out << " ^ ";
            break;
        case AST_OP_LOG_OR:
            _out << " or ";
            break;
        case AST_OP_LOG_AND:
            _out << " and ";
            break;
        case AST_OP_LEFT_SHIFT:
            _out << " << ";
            break;
        case AST_OP_RIGHT_SHIFT:
            _out << " >> ";
            break;
        default:
            printf("Unknown binary op %d\n", op->op_type());
            return op;
    }

    _out << "(";
    visit(op->right());
    _out << ")";
    return op;
}

Node* CodeGen::visit(ListNode* n) {
    auto it = n->node_list()->begin();
    if (it == n->node_list()->end()) {
        return n;
    }

    visit(*it);
    it++;

    for (; it != n->node_list()->end(); it++) {
        _out << ", ";
        visit(*it);
    }
    return n;
}

Node* CodeGen::visit(StmtsNode* n) {    
    for (auto it = n->node_list()->begin(); it != n->node_list()->end(); it++) {
        indent();
        visit(*it);
        _out << "\n";
    }
    return n;
}

unsigned short CodeGen::put_consts_if_absent(Constant& c) {
    auto it = _const_map.find(c);

    int index = -1;
    if (it == _const_map.end()) {
        index = _const_map.size();
        _const_map[c] = index;
        _consts.push_back(c);
    }
    else {
        index = it->second;
    }

    return (unsigned short)index;
}

Node* CodeGen::visit(ConstInt* n) {
    _out << n->value();
    return n;
}

Node* CodeGen::visit(ConstDouble* n) {
    _out << n->value();
    return n;
}

Node* CodeGen::visit(ConstString* n) {
    _out << "\"" << n->_value << "\"";
    return n;
}

Node* CodeGen::visit(ConstArray* n) {
    _out << "[";
    for (auto it = n->arr()->begin(); it != n->arr()->end(); it++) {
        visit(*it);
        if (it != n->arr()->end() - 1) {
            _out << ", ";
        }
    }
    _out << "]";
    return n;
}

Node* CodeGen::visit(BuiltinConst* n) {
    Constant* c = nullptr;
    unsigned short index = 0;
    switch (n->_builtin_id) {
    case BuiltinConst::UNIT:
        c = new Constant(VT_UNIT);
        index = put_consts_if_absent(*c);
        add_op(LOAD_CONST, index);
        break;
    case BuiltinConst::TRUE:
        c = new Constant(VT_TRUE);
        index = put_consts_if_absent(*c);
        add_op(LOAD_CONST, index);
        break;
    case BuiltinConst::FALSE:
        c = new Constant(VT_FALSE);
        index = put_consts_if_absent(*c);
        put_consts_if_absent(*c);
        break;
    default:
        cerr << "unknown const: " << n->_builtin_id;
    }
    return n;
}

Node* CodeGen::visit(ImportNode* n) {
    _out << "from " << n->name() << " import *\n";
    
    return n;
}

Node* CodeGen::visit(ImplementNode* n) {
    // do nothing
    return n;
}

Node* CodeGen::visit(IfNode* n) {
    _out << "if (";
    visit(n->cond());
    _out << "):\n";

    _indent += 1;
    visit(n->then_part());
    _indent -= 1;

    if (n->else_part() != NULL) {
        indent();
        _out << "else:\n";

        _indent += 1;
        visit(n->else_part());
        _indent -= 1;
    }
    return n;
}

Node* CodeGen::visit(WhileNode* n) {
    _out << "while (";
    visit(n->cond());
    _out << "):\n";

    _indent += 1;
    visit(n->body());
    _indent -= 1;
    return n;
}

Node* CodeGen::visit(CmpNode* n) {
    visit(n->left());
    switch (n->cmp_op()) {
    case CmpNode::CMP_LT:
        _out << " < ";
        break;
    case CmpNode::CMP_EQU:
        _out << " == ";
        break;
    case CmpNode::CMP_GT:
        _out << " > ";
        break;
    case CmpNode::CMP_GE:
        _out << " >= ";
        break;
    case CmpNode::CMP_LE:
        _out << " <= ";
        break;
    case CmpNode::CMP_NE:
        _out << " != ";
        break;
    }
    visit(n->right());
    return n;
}

Node* CodeGen::visit(LogicNotNode* n) {
    _out << "not ";
    visit(n->value());
    return n;
}

Node* CodeGen::visit(AssignNode* n) {
    _context = Context::Store;
    visit(n->left());
    _out << " = ";
    _context = Context::Load;
    visit(n->right());
    return n;
}

Node* CodeGen::visit(VarNode* n) {
    if (n->name() == "List") {
        _out << "list";
        return n;
    }

    _out << n->name();
    return n;
}

Node* CodeGen::visit(ParamDefNode* n) {
    _out << n->name();
    return n;
}

Node* CodeGen::visit(VarDefNode* n) {
    if (!n->init_value()) {
        return n;
    }

    int index = _sym_table->get_name_index(n->name());
    if (index >= 0) {
        _out << n->name() << " = ";
        visit(n->init_value());
        return n;
    }

    index = _sym_table->get_var_name_index(n->name());
    if (index >= 0) {
        return n;
    }

    index = _sym_table->get_local_index(n->name());
    if (index >= 0) {
        _out << n->name() << " = ";
        visit(n->init_value());
        return n;
    }

    index = _sym_table->get_cell_var_index(n->name());
    if (index >= 0) {
        _out << n->name() << " = ";
        visit(n->init_value());
        return n;
    }

    // var def should not be free var.
    return n;
}

Node* CodeGen::visit(CallNode* n) {
    Type* callee_type = n->callee()->get_type(type_table());    
    if (callee_type->is_record()) {
        RecordType* rt = dynamic_cast<RecordType*>(callee_type);
        _out << rt->name();
    }
    else {
        visit(n->callee());
    }
    _out << "(";
    visit(n->args());
    _out << ")";
    return n;
}

Node* CodeGen::visit(SubscrNode* n) {
    // For generic instantiation.
    if (n->name()->get_type(type_table())->is_poly() ||
        n->name()->get_type(type_table())->is_type_func()) {
        visit(n->name());
        return n;
    }

    visit(n->name());
    _out << "[";
    visit(n->index());
    _out << "]";
    return n;
}

Node* CodeGen::visit(DotNode* n) {
    visit(n->name());
    _out << ".";
    visit(n->field());
    return n;
}

Node* CodeGen::visit(SequenceNode* n) {
    _out << "HiRange(";
    visit(n->left());
    _out << ", ";
    visit(n->right());
    _out << ")";
    return n;
}

Node* CodeGen::visit(TupleNode* n) {
    _out << "(";
    visit(n->args());
    _out << ")";
    return n;
}

Node* CodeGen::visit(FuncDefNode* n) {
    _sym_table = n->sym_table();
    _out << "def ";
    _out << n->name() << "(";
    visit(n->args());
    _out << "):\n";

    _indent += 1;
    if (n->body()) {
        if (_sym_table->freevars().size() > 0) {
            for (auto it = _sym_table->freevars().begin();
                it != _sym_table->freevars().end(); it++) {
                indent();
                _out << "nonlocal " << *it << "\n";
            }
        }
        visit(n->body());
    }
    else {
        indent();
        _out << "pass\n";
    }
    _indent -= 1;
    _sym_table = _sym_table->parent();
    return n;
}

Node* CodeGen::visit(RecordDefNode* n) {
    _sym_table = n->sym_table();
    _out << "class " << n->name() << ":\n";
    _indent += 1;

    indent();
    _out << "def __init__(self";
    for (auto it = n->fields()->node_list()->begin();
        it != n->fields()->node_list()->end(); it++) {
        _out << ", " << dynamic_cast<VarDefNode*>(*it)->name();
    }
    _out << "):\n";

    _indent += 1;
    if (n->fields()->node_list()->size() > 0) {
        for (auto it = n->fields()->node_list()->begin();
            it != n->fields()->node_list()->end(); it++) {
            indent();
            string name = dynamic_cast<VarDefNode*>(*it)->name();
            _out << "self." << name << " = " << name << "\n";
        }
    }
    else {
        indent();
        _out << "pass\n";
    }

    _indent -= 1;

    _out << "\n";

    indent();
    _out << "def __unapply__(self):\n";
    _indent += 1;
    indent();
    _out << "return (";
    for (auto it = n->fields()->node_list()->begin();
        it != n->fields()->node_list()->end(); it++) {
        string name = dynamic_cast<VarDefNode*>(*it)->name();
        _out << "self." << name << ", "; 
    }
    _out << ")\n";
    _indent -= 1;

    _out << "\n";

    if (n->methods() && n->methods()->size() > 0) {
        for (auto it = n->methods()->node_list()->begin();
            it != n->methods()->node_list()->end(); it++) {
            indent();
            visit(*it);
            _out << "\n";
        }
    }

    // end of class definition.
    _indent -= 1;
    _out << "\n";
    _sym_table = _sym_table->parent();
    return n;
}

Node* CodeGen::visit(TraitDefNode* n) {
    _sym_table = n->sym_table();
    _out << "class " << n->name() << ":\n";
    _indent += 1;
    if (n->methods()) {
        visit(n->methods());
    }
    else {
        indent();
        _out << "pass\n";
    }
    _indent -= 1;
    _sym_table = _sym_table->parent();
    return n;
}

Node* CodeGen::visit(LambdaDef* n) {
    _sym_table = n->sym_table();
    _out << "lambda ";
    visit(n->param());
    _out << ":";

    _indent += 1;
    bool old_value = _is_lambda;
    _is_lambda = true;
    visit(n->body());
    _is_lambda = old_value;
    _indent -= 1;
    _sym_table = _sym_table->parent();
    return n;
}

Node* CodeGen::visit(TypeVarNode* n) {
    // do nothing.
    return n;
}

Node* CodeGen::visit(ArrowTypeNode* n) {
    // do nothing.
    return n;
}

Node* CodeGen::visit(TypeConstInt* n) {
    // do nothing
    return n;
}

Node* CodeGen::visit(ArrayTypeNode* n) {
    // do nothing.
    return n;
}

Node* CodeGen::visit(TypeVarDefNode* n) {
    // do nothing.
    return n;
}

Node* CodeGen::visit(TypeFuncDefNode* n) {
    // do nothing.
    return n;
}

Node* CodeGen::visit(TypeCallNode* n) {
    // do nothing.
    return n;
}

Node* CodeGen::visit(TypeTupleNode* n) {
    // do nothing.
    return n;
}

CodeObject* CodeGen::make_code_object() {
    return new CodeObject(_arg_count, _nlocals, _stack_size, _flags, _insts, _consts, 
        _sym_table->names(), _sym_table->var_names(), _sym_table->freevars(),
        _sym_table->cellvars());
}

void Disassembler::print_detail(unsigned char c, unsigned short param) {
    if (c == COMPARE_OP) {
        switch(param) {
        case LESS:
            printf("(<)");
            break;
        case EQUAL:
            printf("(==)");
            break;
        case LESS_EQUAL:
            printf("(<=)");
            break;
        default:
            printf("(unknown compare operator : %d)", param);
        }
    }
}

void Disassembler::dis(vector<unsigned char>& insts) {
    int pc = 0;
    for (auto it = insts.begin(); it != insts.end(); it++) {
        unsigned char c = *it;
        if (c >= HAVE_ARGUMENT) {
            printf("%3d ", pc);
            printf("%s\t\t", bytecodes[c].c_str());
            unsigned short param = ((unsigned short)*(++it)) |
                    (((unsigned short)*(++it)) << 8);
            printf("%d", param);
            print_detail(c, param);
            printf("\n");
            pc += 3;
        }
        else {
            printf("%3d ", pc);
            printf("%s\n", bytecodes[c].c_str());
            pc += 1;
        }
    }
}

